# -*-coding: utf-8 -*-
"""
    @Theme   : 
    @Time    : 2024/7/16 9:54
    @Author  : YamFish
    @Url     : https://www.zhihu.com/people/yamfish
"""
import math
import time
import json
import hashlib
from contextlib import contextmanager
from struct import pack, unpack
from typing import Generator, Tuple, List

from qanda.utils import login_expire


class RedisUtil:
    def __init__(self, app, client):
        self.key_session = app.config['REDIS_KEY_SESSION']
        self.key_session_expire = app.config['REDIS_KEY_SESSION_EXPIRE']
        self.key_user_bloom = app.config['REDIS_KEY_USER_BLOOM']
        self.key_user_answer = app.config['REDIS_KEY_USER_ANSWER']
        self.key_user_vote = app.config['REDIS_KEY_USER_VOTE']
        self.key_question_vote = app.config['REDIS_KEY_QUESTION_VOTE']
        self.key_answer_vote = app.config['REDIS_KEY_ANSWER_VOTE']
        self.client = client

    def save_session(self, session_id, session_info):
        session_json = json.dumps(session_info)
        self.client.hash_set(self.key_session, session_id, session_json)
        self.client.sorted_set_add(self.key_session_expire, session_id, session_info['expire_time'])

    def delete_session(self, session_id):
        self.client.delete_hash_key(self.key_session, session_id)
        self.client.sorted_set_remove(self.key_session_expire, session_id)

    def fetch_session(self, session_id):
        if not session_id:
            return {}
        session_json = self.client.hash_get(self.key_session, session_id)
        if not session_json:
            return {}
        session_data = json.loads(session_json)
        if login_expire(session_data=session_data):
            return {}
        return session_data

    @staticmethod
    def md5_filed(field: str) -> str:
        return hashlib.md5(field.encode('utf-8')).hexdigest()

    def check_user_answer_question(self, user, question_id):
        return self.client.is_hash_field_exists(self.key_user_answer, self.md5_filed(user + question_id))

    def set_answer_flag(self, question_id, user):
        self.client.hash_set(self.key_user_answer, self.md5_filed(user + question_id), '1')

    def set_user_vote(self, question_or_answer_id, user, value: str):
        if value not in ('-1', '0', '1'):
            raise ValueError(f'Vote str value {value} should be -1, 0 or 1.')
        self.client.hash_set(self.key_user_vote, self.md5_filed(user + question_or_answer_id), value)

    def get_user_vote(self, question_or_answer_id, user):
        return self.client.hash_get(self.key_user_vote, self.md5_filed(user + question_or_answer_id))

    def set_bit_value(self, offset_list: Generator) -> bool:
        for offset in offset_list:
            self.client.set_bit(self.key_user_bloom, offset, 1)
        return True

    def is_all_bit_1(self, offset_list: Generator) -> bool:
        return self.client.is_bit_list_set(self.key_user_bloom, offset_list)

    def increase_vote_score(self, vote_data, user):
        if vote_data['doc_type'] == 'question':
            redis_key = self.key_question_vote
        else:
            redis_key = self.key_answer_vote.format(vote_data['question_id'])
        question_or_answer_id = vote_data['doc_id']
        if vote_data['value'] == 'vote_up':
            vote_num = self.get_user_vote(question_or_answer_id, user)
            if not vote_num or vote_num != b'1':
                self.set_user_vote(question_or_answer_id, user, '1')
                self.client.sorted_set_increase_score(redis_key, question_or_answer_id, 1)
            else:
                raise RuntimeError(f"User {user} can't vote up again!")
        else:
            vote_num = self.get_user_vote(question_or_answer_id, user)
            if not vote_num or vote_num != b'-1':
                self.set_user_vote(question_or_answer_id, user, '-1')
                self.client.sorted_set_increase_score(redis_key, question_or_answer_id, -1)
            else:
                raise RuntimeError(f"User {user} can't vote down again!")

    def get_doc_rank_range(self, doc_type, start, offset, question_id=None) -> List[Tuple[bytes, int]]:
        redis_key = self.key_question_vote if doc_type == 'question' else self.key_answer_vote.format(question_id)
        doc_id_score_list = self.client.sorted_set_get_by_rank(redis_key, start, offset - 1)
        return doc_id_score_list

    def add_question_vote_set(self, doc_id: str):
        self.client.sorted_set_add(self.key_question_vote, doc_id, 0)

    def add_answer_vote_set(self, question_id, answer_id):
        redis_key = self.key_answer_vote.format(question_id)
        self.client.sorted_set_add(redis_key, answer_id, 0)


class BloomFilter:
    def __init__(self, app, redis_util: RedisUtil):
        self.capacity = app.config['BLOOM_CAPACITY']
        self.error_rate = app.config['BLOOM_ERROR_RATE']
        self.num_slices = int(math.ceil(math.log(1.0 / self.error_rate, 2)))
        self.bits_per_slice = int(math.ceil((self.capacity * abs(math.log(self.error_rate))) /
                                            (self.num_slices * (math.log(2) ** 2))))
        self.hash_funcs = self.make_hash_funcs(self.num_slices, self.bits_per_slice)
        self.redis_util = redis_util

    @staticmethod
    def make_hash_funcs(num_slices, num_bits):
        if num_bits >= (1 << 31):
            fmt_code, chunk_size = 'Q', 8
        elif num_bits >= (1 << 15):
            fmt_code, chunk_size = 'I', 4
        else:
            fmt_code, chunk_size = 'H', 2
        total_hash_bits = 8 * num_slices * chunk_size
        if total_hash_bits > 384:
            hash_func = hashlib.sha512
        elif total_hash_bits > 256:
            hash_func = hashlib.sha384
        elif total_hash_bits > 160:
            hash_func = hashlib.sha256
        elif total_hash_bits > 128:
            hash_func = hashlib.sha1
        else:
            hash_func = hashlib.md5
        fmt = fmt_code * (hash_func().digest_size // chunk_size)
        num_salts, extra = divmod(num_slices, len(fmt))
        if extra:
            num_salts += 1
        salts = tuple(hash_func(hash_func(pack('I', i)).digest()) for i in range(num_salts))

        def _make_hash_funcs(key):
            if isinstance(key, str):
                key = key.encode('utf-8')
            else:
                key = str(key).encode('utf-8')
            i = 0
            for salt in salts:
                h = salt.copy()
                h.update(key)
                for uint in unpack(fmt, h.digest()):
                    yield uint % num_bits
                    i += 1
                    if i >= num_slices:
                        return

        return _make_hash_funcs

    def check_duplicate(self, user):
        return self.redis_util.is_all_bit_1(self.hash_funcs(user))

    def set_key(self, user):
        self.redis_util.set_bit_value(self.hash_funcs(user))

    @contextmanager
    def bloom_lock(self, key, acquire_timeout=30, lock_timeout=30):
        key = f'qa_system:bloom:lock:{key}'
        end = time.time() + acquire_timeout
        while time.time() < end:
            if self.redis_util.client.client.set(key, 1, ex=lock_timeout, nx=True):
                try:
                    yield
                finally:
                    self.redis_util.client.client.delete(key)
                    return
            time.sleep(0.001)
        raise RuntimeError("acquire bloom lock timeout.")
